import json
import asyncio
from datetime import datetime
from typing import Optional
from openai import OpenAI
from sqlalchemy.ext.asyncio import AsyncSession

from app import crud, schemas
from app.core.config import settings
from app.core.database import AsyncSessionLocal
from app.core.logging_config import get_logger

logger = get_logger(__name__)


async def extract_medical_info_async(
    content: str,
    record_id: str,
    image_url: Optional[str] = None
) -> None:
    """
    异步提取医疗报告信息并更新医疗档案记录

    Args:
        content: 文件解析后的文本内容
        record_id: 医疗档案记录ID
        image_url: 图片URL（用于图片类型的医疗档案）
    """
    # 创建新的数据库会话和AI客户端
    async with AsyncSessionLocal() as db:
        llm_client = OpenAI(
            api_key=settings.OPENAI_API_KEY,
            base_url=settings.OPENAI_BASE_URL
        )
        
        try:
            # 更新状态为处理中
            # 添加重试机制，以防主事务还未提交
            retry_count = 0
            max_retries = 3
            record_obj = None
            
            while retry_count < max_retries:
                record_obj = await crud.crud_medical_record.get(db, id=record_id)
                if record_obj:
                    break
                retry_count += 1
                if retry_count < max_retries:
                    # 等待一秒后重试
                    await asyncio.sleep(1)
            
            if record_obj:
                await crud.crud_medical_record.update(
                    db, 
                    db_obj=record_obj, 
                    obj_in={"extraction_status": "processing"}
                )
                await db.commit()
            else:
                logger.error(f"Medical record {record_id} not found after {max_retries} retries")
                return

            # 构造提取医疗信息的提示词和消息
            if image_url:
                # 图片分析
                extraction_prompt = """
                请分析这张医疗报告图片，提取关键信息并以JSON格式返回。如果某些信息无法确定，请返回null。

                请提取以下信息：
                1. disease_name: 病症名称（主要诊断或疾病名称）
                2. report_category: 报告分类（如：检验报告、影像报告、病理报告、体检报告等）
                3. report_name: 报告名称（具体的检查项目名称）
                4. report_date: 报告时间（格式：YYYY-MM-DD，如果只有年月则为YYYY-MM-01）
                5. hospital_name: 就诊医院（医院或医疗机构名称）

                请严格按照以下JSON格式返回，不要包含任何其他文字：
                {
                  "disease_name": "string or null",
                  "report_category": "string or null",
                  "report_name": "string or null",
                  "report_date": "YYYY-MM-DD or null",
                  "hospital_name": "string or null"
                }
                """

                messages = [
                    {"role": "system", "content": "你是一个专业的医疗报告分析助手，擅长从医疗报告中提取结构化信息。请严格按照要求的JSON格式返回结果。"},
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": extraction_prompt},
                            {"type": "image_url", "image_url": {"url": image_url}}
                        ]
                    }
                ]
            else:
                # 文本分析
                extraction_prompt = f"""
                请分析以下医疗报告内容，提取关键信息并以JSON格式返回。如果某些信息无法确定，请返回null。

                医疗报告内容：
                {content}

                请提取以下信息：
                1. disease_name: 病症名称（主要诊断或疾病名称）
                2. report_category: 报告分类（如：检验报告、影像报告、病理报告、体检报告等）
                3. report_name: 报告名称（具体的检查项目名称）
                4. report_date: 报告时间（格式：YYYY-MM-DD，如果只有年月则为YYYY-MM-01）
                5. hospital_name: 就诊医院（医院或医疗机构名称）

                请严格按照以下JSON格式返回，不要包含任何其他文字：
                """+"""
                {
                  "disease_name": "string or null",
                  "report_category": "string or null",
                  "report_name": "string or null",
                  "report_date": "YYYY-MM-DD or null",
                  "hospital_name": "string or null"
                }
                """

                messages = [
                    {"role": "system", "content": "你是一个专业的医疗报告分析助手，擅长从医疗报告中提取结构化信息。请严格按照要求的JSON格式返回结果。"},
                    {"role": "user", "content": extraction_prompt}
                ]

            # 调用AI模型提取信息
            response = llm_client.chat.completions.create(
                model=settings.CHAT_VISION_MODEL,
                messages=messages,
                temperature=0.1,  # 降低温度以获得更一致的结果
                response_format={"type": "json_object"}
            )

            # 解析AI返回的JSON
            ai_response = response.choices[0].message.content.strip()

            # 去除可能的markdown代码块标记
            if ai_response.startswith('```'):
                ai_response = ai_response.split('```')[1]
                if ai_response.startswith('json'):
                    ai_response = ai_response[4:]
        
            try:
                extracted_data = json.loads(ai_response)
            except json.JSONDecodeError:
                logger.error(f"Failed to parse AI response as JSON: {ai_response}")
                extracted_data = {}

            # 验证和处理提取的数据
            medical_info = schemas.MedicalReportExtraction(**extracted_data)

            # 处理日期格式
            report_date = None
            if medical_info.report_date:
                try:
                    report_date = datetime.strptime(medical_info.report_date, "%Y-%m-%d")
                except ValueError:
                    # 尝试其他日期格式
                    try:
                        report_date = datetime.strptime(medical_info.report_date, "%Y-%m")
                        report_date = report_date.replace(day=1)
                    except ValueError:
                        logger.warning(f"Failed to parse date: {medical_info.report_date}")
                        report_date = None

            # 更新医疗档案记录
            update_data = {
                "disease_name": medical_info.disease_name,
                "report_category": medical_info.report_category,
                "report_name": medical_info.report_name,
                "report_date": report_date,
                "hospital_name": medical_info.hospital_name,
                "extraction_status": "completed"
            }

            record_obj = await crud.crud_medical_record.get(db, id=record_id)
            if record_obj:
                await crud.crud_medical_record.update(db, db_obj=record_obj, obj_in=update_data)
                await db.commit()

            logger.info(f"Successfully extracted medical info for record {record_id}")

        except Exception as e:
            logger.error(f"Error extracting medical info for record {record_id}: {str(e)}")

            # 更新状态为失败
            try:
                record_obj = await crud.crud_medical_record.get(db, id=record_id)
                if record_obj:
                    await crud.crud_medical_record.update(
                        db,
                        db_obj=record_obj,
                        obj_in={"extraction_status": "failed"}
                    )
                    await db.commit()
            except Exception as update_error:
                logger.error(f"Failed to update extraction status: {update_error}")